/*
Copyright (C) 2018-2019 de4dot@gmail.com

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

mod call_instr;
mod ip_relmem_instr;
mod jcc_instr;
mod jmp_instr;
mod simple_br_instr;
mod simple_instr;
mod xbegin_instr;

use self::call_instr::CallInstr;
use self::ip_relmem_instr::IpRelMemOpInstr;
use self::jcc_instr::JccInstr;
use self::jmp_instr::JmpInstr;
use self::simple_br_instr::SimpleBranchInstr;
use self::simple_instr::SimpleInstr;
use self::xbegin_instr::XbeginInstr;
use super::block::{Block, BlockData};
use super::*;
#[cfg(any(has_alloc, not(feature = "std")))]
use alloc::rc::Rc;
#[cfg(not(feature = "std"))]
use alloc::string::String;
use core::cell::RefCell;
use core::i32;
#[cfg(all(not(has_alloc), feature = "std"))]
use std::rc::Rc;

pub(super) trait Instr {
	fn block(&self) -> Rc<RefCell<Block>>;
	fn size(&self) -> u32;
	fn ip(&self) -> u64;
	fn set_ip(&mut self, new_ip: u64);
	fn orig_ip(&self) -> u64;

	/// Initializes the target address and tries to optimize the instruction
	fn initialize(&mut self, block_encoder: &BlockEncoder);

	/// Returns `true` if the instruction was updated to a shorter instruction, `false` if nothing changed
	fn optimize(&mut self) -> bool;

	fn encode(&mut self, block: &mut Block) -> Result<(ConstantOffsets, bool), String>;
}

#[derive(Default)]
pub(super) struct TargetInstr {
	instruction: Option<Rc<RefCell<Instr>>>,
	address: u64,
	is_owner: bool,
}

impl TargetInstr {
	#[inline]
	pub(super) fn new_instr(instruction: Rc<RefCell<Instr>>) -> Self {
		Self { instruction: Some(Rc::clone(&instruction)), address: 0, is_owner: false }
	}

	#[inline]
	pub(super) fn new_address(address: u64) -> Self {
		Self { instruction: None, address, is_owner: false }
	}

	#[inline]
	pub(super) fn new_owner() -> Self {
		Self { instruction: None, address: 0, is_owner: true }
	}

	fn is_in_block(&self, block: Rc<RefCell<Block>>) -> bool {
		if self.is_owner {
			// The owner checks if the input block is part of its block, so return true
			true
		} else if let Some(ref instr) = self.instruction {
			Rc::ptr_eq(&instr.borrow().block(), &block)
		} else {
			false
		}
	}

	fn address(&self, owner: &Instr) -> u64 {
		if self.is_owner {
			owner.ip()
		} else {
			match self.instruction {
				Some(ref instr) => instr.borrow().ip(),
				None => self.address,
			}
		}
	}
}

pub(super) struct InstrUtils;
impl InstrUtils {
	// 6 = FF 15 XXXXXXXX = call qword ptr [rip+mem_target]
	pub(self) const CALL_OR_JMP_POINTER_DATA_INSTRUCTION_SIZE64: u32 = 6;

	#[cfg(any(feature = "gas", feature = "intel", feature = "masm", feature = "nasm"))]
	pub(self) fn create_error_message(error_message: &str, instruction: &Instruction) -> String {
		format!("{} : 0x{:X} {}", error_message, instruction.ip(), instruction)
	}

	#[cfg(not(any(feature = "gas", feature = "intel", feature = "masm", feature = "nasm")))]
	pub(self) fn create_error_message(error_message: &str, instruction: &Instruction) -> String {
		format!("{} : 0x{:X}", error_message, instruction.ip())
	}

	pub(super) fn create<'a, 'b>(block_encoder: &'a mut BlockEncoder, block: Rc<RefCell<Block>>, instruction: &'b Instruction) -> Rc<RefCell<Instr>> {
		#[cfg_attr(feature = "cargo-fmt", rustfmt::skip)]
		match instruction.code() {
			// GENERATOR-BEGIN: JccInstr
			// ⚠️This was generated by GENERATOR!🦹‍♂️
			Code::Jo_rel8_16
			| Code::Jo_rel8_32
			| Code::Jo_rel8_64
			| Code::Jno_rel8_16
			| Code::Jno_rel8_32
			| Code::Jno_rel8_64
			| Code::Jb_rel8_16
			| Code::Jb_rel8_32
			| Code::Jb_rel8_64
			| Code::Jae_rel8_16
			| Code::Jae_rel8_32
			| Code::Jae_rel8_64
			| Code::Je_rel8_16
			| Code::Je_rel8_32
			| Code::Je_rel8_64
			| Code::Jne_rel8_16
			| Code::Jne_rel8_32
			| Code::Jne_rel8_64
			| Code::Jbe_rel8_16
			| Code::Jbe_rel8_32
			| Code::Jbe_rel8_64
			| Code::Ja_rel8_16
			| Code::Ja_rel8_32
			| Code::Ja_rel8_64
			| Code::Js_rel8_16
			| Code::Js_rel8_32
			| Code::Js_rel8_64
			| Code::Jns_rel8_16
			| Code::Jns_rel8_32
			| Code::Jns_rel8_64
			| Code::Jp_rel8_16
			| Code::Jp_rel8_32
			| Code::Jp_rel8_64
			| Code::Jnp_rel8_16
			| Code::Jnp_rel8_32
			| Code::Jnp_rel8_64
			| Code::Jl_rel8_16
			| Code::Jl_rel8_32
			| Code::Jl_rel8_64
			| Code::Jge_rel8_16
			| Code::Jge_rel8_32
			| Code::Jge_rel8_64
			| Code::Jle_rel8_16
			| Code::Jle_rel8_32
			| Code::Jle_rel8_64
			| Code::Jg_rel8_16
			| Code::Jg_rel8_32
			| Code::Jg_rel8_64
			| Code::Jo_rel16
			| Code::Jo_rel32_32
			| Code::Jo_rel32_64
			| Code::Jno_rel16
			| Code::Jno_rel32_32
			| Code::Jno_rel32_64
			| Code::Jb_rel16
			| Code::Jb_rel32_32
			| Code::Jb_rel32_64
			| Code::Jae_rel16
			| Code::Jae_rel32_32
			| Code::Jae_rel32_64
			| Code::Je_rel16
			| Code::Je_rel32_32
			| Code::Je_rel32_64
			| Code::Jne_rel16
			| Code::Jne_rel32_32
			| Code::Jne_rel32_64
			| Code::Jbe_rel16
			| Code::Jbe_rel32_32
			| Code::Jbe_rel32_64
			| Code::Ja_rel16
			| Code::Ja_rel32_32
			| Code::Ja_rel32_64
			| Code::Js_rel16
			| Code::Js_rel32_32
			| Code::Js_rel32_64
			| Code::Jns_rel16
			| Code::Jns_rel32_32
			| Code::Jns_rel32_64
			| Code::Jp_rel16
			| Code::Jp_rel32_32
			| Code::Jp_rel32_64
			| Code::Jnp_rel16
			| Code::Jnp_rel32_32
			| Code::Jnp_rel32_64
			| Code::Jl_rel16
			| Code::Jl_rel32_32
			| Code::Jl_rel32_64
			| Code::Jge_rel16
			| Code::Jge_rel32_32
			| Code::Jge_rel32_64
			| Code::Jle_rel16
			| Code::Jle_rel32_32
			| Code::Jle_rel32_64
			| Code::Jg_rel16
			| Code::Jg_rel32_32
			| Code::Jg_rel32_64
			=> return Rc::new(RefCell::new(JccInstr::new(block_encoder, block, instruction))),
			// GENERATOR-END: JccInstr

			// GENERATOR-BEGIN: SimpleBranchInstr
			// ⚠️This was generated by GENERATOR!🦹‍♂️
			Code::Loopne_rel8_16_CX
			| Code::Loopne_rel8_32_CX
			| Code::Loopne_rel8_16_ECX
			| Code::Loopne_rel8_32_ECX
			| Code::Loopne_rel8_64_ECX
			| Code::Loopne_rel8_16_RCX
			| Code::Loopne_rel8_64_RCX
			| Code::Loope_rel8_16_CX
			| Code::Loope_rel8_32_CX
			| Code::Loope_rel8_16_ECX
			| Code::Loope_rel8_32_ECX
			| Code::Loope_rel8_64_ECX
			| Code::Loope_rel8_16_RCX
			| Code::Loope_rel8_64_RCX
			| Code::Loop_rel8_16_CX
			| Code::Loop_rel8_32_CX
			| Code::Loop_rel8_16_ECX
			| Code::Loop_rel8_32_ECX
			| Code::Loop_rel8_64_ECX
			| Code::Loop_rel8_16_RCX
			| Code::Loop_rel8_64_RCX
			| Code::Jcxz_rel8_16
			| Code::Jcxz_rel8_32
			| Code::Jecxz_rel8_16
			| Code::Jecxz_rel8_32
			| Code::Jecxz_rel8_64
			| Code::Jrcxz_rel8_16
			| Code::Jrcxz_rel8_64
			=> return Rc::new(RefCell::new(SimpleBranchInstr::new(block_encoder, block, instruction))),
			// GENERATOR-END: SimpleBranchInstr

			// GENERATOR-BEGIN: CallInstr
			// ⚠️This was generated by GENERATOR!🦹‍♂️
			Code::Call_rel16
			| Code::Call_rel32_32
			| Code::Call_rel32_64
			=> return Rc::new(RefCell::new(CallInstr::new(block_encoder, block, instruction))),
			// GENERATOR-END: CallInstr

			// GENERATOR-BEGIN: JmpInstr
			// ⚠️This was generated by GENERATOR!🦹‍♂️
			Code::Jmp_rel16
			| Code::Jmp_rel32_32
			| Code::Jmp_rel32_64
			| Code::Jmp_rel8_16
			| Code::Jmp_rel8_32
			| Code::Jmp_rel8_64
			=> return Rc::new(RefCell::new(JmpInstr::new(block_encoder, block, instruction))),
			// GENERATOR-END: JmpInstr

			// GENERATOR-BEGIN: XbeginInstr
			// ⚠️This was generated by GENERATOR!🦹‍♂️
			Code::Xbegin_rel16
			| Code::Xbegin_rel32
			=> return Rc::new(RefCell::new(XbeginInstr::new(block_encoder, block, instruction))),
			// GENERATOR-END: XbeginInstr

			_ => {},
		}

		if block_encoder.bitness() == 64 {
			for i in 0..instruction.op_count() {
				if instruction.op_kind(i) == OpKind::Memory {
					if instruction.is_ip_rel_memory_operand() {
						return Rc::new(RefCell::new(IpRelMemOpInstr::new(block_encoder, block, instruction)));
					}
					break;
				}
			}
		}

		Rc::new(RefCell::new(SimpleInstr::new(block_encoder, block, instruction)))
	}

	fn encode_branch_to_pointer_data(
		block: &mut Block, is_call: bool, ip: u64, pointer_data: Rc<RefCell<BlockData>>, min_size: u32,
	) -> Result<u32, String> {
		if min_size > i32::MAX as u32 {
			return Err(String::from("Internal error: min_size > i32::MAX"));
		}

		let mut instr = Instruction::default();
		instr.set_op0_kind(OpKind::Memory);
		instr.set_memory_displ_size(block.encoder.bitness() / 8);
		let reloc_kind;
		match block.encoder.bitness() {
			64 => {
				instr.set_code(if is_call { Code::Call_rm64 } else { Code::Jmp_rm64 });
				instr.set_memory_base(Register::RIP);
				let next_rip = ip.wrapping_add(Self::CALL_OR_JMP_POINTER_DATA_INSTRUCTION_SIZE64 as u64);
				instr.set_next_ip(next_rip);
				let diff = pointer_data.borrow().address().wrapping_sub(next_rip) as i64;
				instr.set_memory_displacement(diff as u32);
				if !(i32::MIN as i64 <= diff && diff <= i32::MAX as i64) {
					return Err(String::from("Block is too big"));
				}
				reloc_kind = RelocKind::Offset64;
			}

			_ => unreachable!(),
		}

		let mut size = match block.encoder.encode(&instr, ip) {
			Ok(len) => len,
			Err(err) => return Err(err),
		} as u32;
		if block.can_add_reloc_infos() && reloc_kind != RelocKind::Offset64 {
			let co = block.encoder.get_constant_offsets();
			if !co.has_displacement() {
				return Err(String::from("Internal error: no displ"));
			}
			block.add_reloc_info(RelocInfo::new(reloc_kind, ip.wrapping_add(co.displacement_offset() as u64)));
		}
		while size < min_size {
			size += 1;
			block.write_byte(0x90);
		}
		Ok(size)
	}
}
